import tensorflow as tf
import tensor.top_k_accuracy as tka

prob = tf.constant([[0.1, 0.2, 0.7], [0.2, 0.7, 0.1]])
target = tf.constant([2, 0])

t_k = tf.math.top_k(prob, 3)
print(t_k.indices)
print(t_k.values)

t_k_p = tf.transpose(t_k.indices, [1, 0])
print(t_k_p)

target_bd = tf.broadcast_to(target, [3, 2])
print(target_bd)

correct = tf.equal(t_k_p, target_bd)
print(correct)

print(correct[:2])
correct_2 = tf.reshape(correct[:2], [-1])
print(correct_2)
correct_float_2 = tf.cast(correct[:2], dtype=tf.float32)
print(correct_float_2)
print(tf.reduce_sum(correct_float_2))
accuracy = float(tf.reduce_sum(correct_float_2) / 2)
print(accuracy)

accuracy = tka.top_k_accuracy(prob, target, top_k=(1, 2, 3))
print('top accuracy:', accuracy)
